import cv2
from numba import jit
import math
import numpy as np
import time
import subprocess


@jit(nopython=True)
def cylindrical_projection(img, f):
    rows = img.shape[0]
    cols = img.shape[1]

    blank = np.zeros_like(img)
    center_x = int(cols / 2)
    center_y = int(rows / 2)

    for y in range(rows):
        for x in range(cols):
            theta = math.atan((x - center_x) / f)
            point_x = int(f * math.tan((x - center_x) / f) + center_x)
            point_y = int((y - center_y) / math.cos(theta) + center_y)

            if point_x >= cols or point_x < 0 or point_y >= rows or point_y < 0:
                pass
            else:
                blank[y, x, :] = img[point_y, point_x, :]
    return blank

@jit(nopython=True)
def image_cut(image,binary_image):
    x = binary_image.shape[0]
    y = binary_image.shape[1]

    edges_x = []
    edges_y = []

    for i in range(x):
        for j in range(y):
            if binary_image[i][j] == 255:
                edges_x.append(i)
                edges_y.append(j)

    left = min(edges_x) + 50  # 左边界
    right = max(edges_x)  # 右边界
    width = right - left - 50  # 宽度

    bottom = min(edges_y)  # 底部
    top = max(edges_y)  # 顶部
    height = top - bottom  # 高度

    pre1_picture = image[left:left + width, bottom:bottom + height]  # 图片截取

    return pre1_picture  # 返回图片数据


def change_size(image):
    b = cv2.threshold(image, 15, 255, cv2.THRESH_BINARY)  # 调整裁剪效果
    binary_image = b[1]  # 二值图--具有三通道
    binary_image = cv2.cvtColor(binary_image, cv2.COLOR_BGR2GRAY) # 改为单通道
    return image_cut(image,binary_image)

def push_rtsp_stream(url, rtsp_url, ratio):
    # 打开摄像头
    cap = cv2.VideoCapture(url)
    if not cap.isOpened():
        print("Could not open camera")
        return

    # 设置视频编解码器参数
    fourcc = cv2.VideoWriter_fourcc(*'H264')
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    frame_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))

    # 创建视频编码器
    writer = cv2.VideoWriter('pipe:', fourcc, fps, frame_size, isColor=True)

    # 创建FFmpeg进程并推流
    cmd = ['ffmpeg', '-y', '-f', 'rawvideo', '-pixel_format', 'bgr24', '-video_size', f"{frame_size[0]}x{frame_size[1]}", '-i', 'pipe:', '-c:v', 'libx264', '-preset', 'ultrafast', '-tune', 'zerolatency', '-f', 'rtsp', rtsp_url]
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE)

    # 循环从摄像头读取帧并推流
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        img = cylindrical_projection(frame, ratio)
        result = change_size(image=img) if cylindrical_projection else frame

        writer.write(result)
        p.stdin.write(result.tobytes())

    # 释放资源
    cap.release()
    writer.release()
    p.stdin.close()
    p.wait()

def start(url, ratio):
    cap = cv2.VideoCapture(url)
    if cap.isOpened():
        print("连接成功!")
    else:
        print("连接失败!")

    # print("fps:", cap.get(cv2.CAP_PROP_FPS))
    size = (int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)))
    print("视频大小：", size)

    pTime = 0
    while (True):
        # Capture frame-by-frame
        ret, frame = cap.read()
        # Our operations on the frame come here
        img = cylindrical_projection(frame, ratio)
        result = change_size(image=img) if cylindrical_projection else frame
        # show fps
        cTime = time.time()
        fps = 1 / (cTime - pTime)
        pTime = cTime
        print("fps:" + str(int(fps)))
        # Display the resulting frame
        img = cv2.resize(img, (800, 600))  # 缩小展示
        cv2.imshow('img', img)
        result = cv2.resize(result, (800, 600))  # 缩小展示
        cv2.imshow('result', result)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
